package ssm.config;

import com.google.common.collect.Lists;
import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;

import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.*;

public class TransactionRangeShardingAlgorithm implements RangeShardingAlgorithm<Date> {

    private static final DateTimeFormatter YYYY_MM = DateTimeFormatter.ofPattern("YYYYMM");
    private static final Integer SIZE = 3;

    private static LocalDateTime dateToLocalDateTime(Date date) {
        Instant instant = date.toInstant();
        ZoneId zoneId = ZoneId.systemDefault();
        // atZone()方法返回在指定时区从此Instant生成的ZonedDateTime。
        return instant.atZone(zoneId).toLocalDateTime();

    }

    private static List<String> getTableNames(LocalDateTime start, LocalDateTime end) {
        List<String> result = Lists.newArrayList();
        String startString = YYYY_MM.format(start);
        //定义：transaction 表只能对近三个月的数据进行操作
        for (int i = 0; i < SIZE; i++) {
            LocalDateTime endTime = end.plusMonths(i);
            String endString = YYYY_MM.format(endTime);
            result.add(endString);
            if (startString.equals(endString)) {
                break;
            }
        }
        return result;
    }

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, RangeShardingValue shardingValue) {
        Range<Date> valueRange = shardingValue.getValueRange();
        Date lowerEnd = valueRange.lowerEndpoint();
        Date upperEnd = valueRange.upperEndpoint();
        Set<String> routTables = new HashSet<>();
        if (lowerEnd != null && upperEnd != null) {
            List<String> rangeNameList = getTableNames(dateToLocalDateTime(lowerEnd), dateToLocalDateTime(upperEnd));
            for (String string : rangeNameList) {
                routTables.add(shardingValue.getLogicTableName() + "_" + string);
            }
        }
        return routTables;
    }
}
